{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNFl8iwnrTusJCiB22sgXKk",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mshumer/ai-journalist/blob/main/Claude_Journalist.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install newspaper3k"
      ],
      "metadata": {
        "id": "KwLlmf_8EikO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7t2ITaBSCj7n"
      },
      "outputs": [],
      "source": [
        "import requests\n",
        "from bs4 import BeautifulSoup\n",
        "import newspaper\n",
        "from newspaper import Article\n",
        "import ast\n",
        "\n",
        "ANTHROPIC_API_KEY = \"YOUR API KEY\"  # Replace with your Anthropic API key\n",
        "SERP_API_KEY = \"YOUR API KEY\"  # Replace with your SERP API key\n",
        "\n",
        "def get_search_terms(topic):\n",
        "    system_prompt = \"You are a world-class journalist. Generate a list of 5 search terms to search for to research and write an article about the topic.\"\n",
        "    messages = [\n",
        "        {\"role\": \"user\", \"content\": f\"Please provide a list of 5 search terms related to '{topic}' for researching and writing an article. Respond with the search terms in a Python-parseable list, separated by commas.\"},\n",
        "    ]\n",
        "    headers = {\n",
        "        \"x-api-key\": ANTHROPIC_API_KEY,\n",
        "        \"anthropic-version\": \"2023-06-01\",\n",
        "        \"content-type\": \"application/json\"\n",
        "    }\n",
        "    data = {\n",
        "        \"model\": 'claude-3-haiku-20240307',\n",
        "        \"max_tokens\": 200,\n",
        "        \"temperature\": 0.5,\n",
        "        \"system\": system_prompt,\n",
        "        \"messages\": messages,\n",
        "    }\n",
        "\n",
        "    response = requests.post(\"https://api.anthropic.com/v1/messages\", headers=headers, json=data)\n",
        "    response_text = response.json()['content'][0]['text']\n",
        "    search_terms = ast.literal_eval(response_text)\n",
        "    return search_terms\n",
        "\n",
        "def get_search_results(search_term):\n",
        "    url = f\"https://serpapi.com/search.json?q={search_term}&api_key={SERP_API_KEY}\"\n",
        "    response = requests.get(url)\n",
        "    data = response.json()\n",
        "    return data['organic_results']\n",
        "\n",
        "def select_relevant_urls(search_results):\n",
        "    system_prompt = \"You are a journalist assistant. From the given search results, select the URLs that seem most relevant and informative for writing an article on the topic.\"\n",
        "    search_results_text = \"\\n\".join([f\"{i+1}. {result['link']}\" for i, result in enumerate(search_results)])\n",
        "    messages = [\n",
        "        {\"role\": \"user\", \"content\": f\"Search Results:\\n{search_results_text}\\n\\nPlease select the numbers of the URLs that seem most relevant and informative for writing an article on the topic. Respond with the numbers in a Python-parseable list, separated by commas.\"},\n",
        "    ]\n",
        "    headers = {\n",
        "        \"x-api-key\": ANTHROPIC_API_KEY,\n",
        "        \"anthropic-version\": \"2023-06-01\",\n",
        "        \"content-type\": \"application/json\"\n",
        "    }\n",
        "    data = {\n",
        "        \"model\": 'claude-3-haiku-20240307',\n",
        "        \"max_tokens\": 200,\n",
        "        \"temperature\": 0.5,\n",
        "        \"system\": system_prompt,\n",
        "        \"messages\": messages,\n",
        "    }\n",
        "    response = requests.post(\"https://api.anthropic.com/v1/messages\", headers=headers, json=data)\n",
        "    response_text = response.json()['content'][0]['text']\n",
        "\n",
        "    numbers = ast.literal_eval(response_text)\n",
        "    relevant_indices = [int(num) - 1 for num in numbers]\n",
        "    relevant_urls = [search_results[i]['link'] for i in relevant_indices]\n",
        "\n",
        "    return relevant_urls\n",
        "\n",
        "def get_article_text(url):\n",
        "    article = Article(url)\n",
        "    article.download()\n",
        "    article.parse()\n",
        "    return article.text\n",
        "\n",
        "def write_article(topic, article_texts):\n",
        "    system_prompt = \"You are a journalist. Write a high-quality, NYT-worthy article on the given topic based on the provided article texts. The article should be well-structured, informative, and engaging.\"\n",
        "    combined_text = \"\\n\\n\".join(article_texts)\n",
        "    messages = [\n",
        "        {\"role\": \"user\", \"content\": f\"Topic: {topic}\\n\\nArticle Texts:\\n{combined_text}\\n\\nPlease write a high-quality, NYT-worthy article on the topic based on the provided article texts. The article should be well-structured, informative, and engaging. Ensure the length is at least as long as a NYT cover story -- at a minimum, 15 paragraphs.\"},\n",
        "    ]\n",
        "    headers = {\n",
        "        \"x-api-key\": ANTHROPIC_API_KEY,\n",
        "        \"anthropic-version\": \"2023-06-01\",\n",
        "        \"content-type\": \"application/json\"\n",
        "    }\n",
        "    data = {\n",
        "        \"model\": 'claude-3-opus-20240229',\n",
        "        \"max_tokens\": 3000,\n",
        "        \"temperature\": 0.5,\n",
        "        \"system\": system_prompt,\n",
        "        \"messages\": messages,\n",
        "    }\n",
        "    response = requests.post(\"https://api.anthropic.com/v1/messages\", headers=headers, json=data)\n",
        "    article = response.json()['content'][0]['text']\n",
        "    return article\n",
        "\n",
        "def edit_article(article):\n",
        "    system_prompt = \"You are an editor. Review the given article and provide suggestions for improvement. Focus on clarity, coherence, and overall quality.\"\n",
        "    messages = [\n",
        "        {\"role\": \"user\", \"content\": f\"Article:\\n{article}\\n\\nPlease review the article and provide suggestions for improvement. Focus on clarity, coherence, and overall quality.\"},\n",
        "    ]\n",
        "    headers = {\n",
        "        \"x-api-key\": ANTHROPIC_API_KEY,\n",
        "        \"anthropic-version\": \"2023-06-01\",\n",
        "        \"content-type\": \"application/json\"\n",
        "    }\n",
        "    data = {\n",
        "        \"model\": 'claude-3-opus-20240229',\n",
        "        \"max_tokens\": 3000,\n",
        "        \"temperature\": 0.5,\n",
        "        \"system\": system_prompt,\n",
        "        \"messages\": messages,\n",
        "    }\n",
        "    response = requests.post(\"https://api.anthropic.com/v1/messages\", headers=headers, json=data)\n",
        "    suggestions = response.json()['content'][0]['text']\n",
        "\n",
        "    system_prompt = \"You are an editor. Rewrite the given article based on the provided suggestions for improvement.\"\n",
        "    messages = [\n",
        "        {\"role\": \"user\", \"content\": f\"Original Article:\\n{article}\\n\\nSuggestions for Improvement:\\n{suggestions}\\n\\nPlease rewrite the article based on the provided suggestions for improvement.\"},\n",
        "    ]\n",
        "    data = {\n",
        "        \"model\": 'claude-3-opus-20240229',\n",
        "        \"max_tokens\": 3000,\n",
        "        \"temperature\": 0.5,\n",
        "        \"system\": system_prompt,\n",
        "        \"messages\": messages,\n",
        "    }\n",
        "    response = requests.post(\"https://api.anthropic.com/v1/messages\", headers=headers, json=data)\n",
        "    edited_article = response.json()['content'][0]['text']\n",
        "    return edited_article\n",
        "\n",
        "# User input\n",
        "topic = input(\"Enter a topic to write about: \")\n",
        "do_edit = input(\"After the initial draft, do you want an automatic edit? This may improve performance, but is slightly unreliable. Answer 'yes' or 'no'.\")\n",
        "\n",
        "# Generate search terms\n",
        "search_terms = get_search_terms(topic)\n",
        "print(f\"\\nSearch Terms for '{topic}':\")\n",
        "print(\", \".join(search_terms))\n",
        "\n",
        "# Perform searches and select relevant URLs\n",
        "relevant_urls = []\n",
        "for term in search_terms:\n",
        "    search_results = get_search_results(term)\n",
        "    urls = select_relevant_urls(search_results)\n",
        "    relevant_urls.extend(urls)\n",
        "\n",
        "print('Relevant URLs to read:', relevant_urls)\n",
        "\n",
        "\n",
        "# Get article text from relevant URLs\n",
        "article_texts = []\n",
        "for url in relevant_urls:\n",
        "  try:\n",
        "    text = get_article_text(url)\n",
        "    if len(text) > 75:\n",
        "      article_texts.append(text)\n",
        "  except:\n",
        "    pass\n",
        "\n",
        "print('Articles to reference:', article_texts)\n",
        "\n",
        "print('\\n\\nWriting article...')\n",
        "# Write the article\n",
        "article = write_article(topic, article_texts)\n",
        "print(\"\\nGenerated Article:\")\n",
        "print(article)\n",
        "\n",
        "if 'y' in do_edit:\n",
        "  # Edit the article\n",
        "  edited_article = edit_article(article)\n",
        "  print(\"\\nEdited Article:\")\n",
        "  print(edited_article)"
      ]
    }
  ]
}